#!/usr/bin/env python
# encoding: utf-8
"""
test_hmm.py

Created by Orestis Markou on 2007-12-12.
Copyright (c) 2007 Orestis Markou. All rights reserved.
"""

import sys
import os

import unittest
from greeklish.numpy_hmm import MyHMM



class HMMTestCase1(unittest.TestCase):
    def setUp(self):
        state_transitions = {
         'sun':{'sun':0.5, 'cloud':0.25, 'rain':0.25}, 
         'cloud':{'sun':0.375, 'cloud':0.125, 'rain':0.375}, 
         'rain':{'sun':0.125, 'cloud':0.675, 'rain':0.375}
                         }     
        start_probabilities = {'sun':0.63, 'cloud':0.17, 'rain':0.20}

        states = ['sun', 'cloud', 'rain']
        observations = ['dry', 'dryish', 'damp', 'soggy']

        observations_prob = {
        'sun':{'dry':0.60, 'dryish':0.20, 'damp':0.15, 'soggy':0.05}, 
        'cloud':{'dry':0.25, 'dryish':0.25, 'damp':0.25, 'soggy':0.25}, 
        'rain':{'dry':0.05, 'dryish':0.10, 'damp':0.35, 'soggy':0.50}
        }

        self.hmm = MyHMM(states, observations, start_probabilities, state_transitions, observations_prob)

        self.obs_sequence = ['dry', 'damp', 'soggy', 'dry','damp','damp','dry','dry']
        self.result = ['sun', 'rain', 'rain', 'cloud', 'rain', 'cloud', 'sun', 'sun']
    
    def tearDown(self):
        """docstring for tearDown"""
        pass
        
    def testOptimizeMethod(self):
        import numpy
        A = numpy.array([
        [0.5, 0.25, 0.25],
        [0.375,0.125,0.375],
        [0.125,0.675,0.375]
        ])
        Pi = numpy.array(
        [0.63,0.17,0.20]
        )
        B = numpy.array([
        [0.60,0.20,0.15,0.05],
        [0.25,0.25,0.25,0.25],
        [0.05,0.10,0.35,0.50]
        ])
        self.hmm._numpy = True
        self.hmm._optimize()
        # start_probabilities = self._opt_numpy_start_prob
        # observations_prob = self._opt_numpy_obs_prob
        # state_transitions = self._opt_numpy_state_trans
        # self.assert_(A.all()==self.hmm._opt_numpy_obs_prob.all())
        self.assert_((A==self.hmm._opt_numpy_state_trans).all(), 
                'A%s\n---\nA!%s'%(A,self.hmm._opt_numpy_state_trans))
        self.assert_((B==self.hmm._opt_numpy_obs_prob).all()        , 
                'B%s\n---\nB!%s'%(B,self.hmm._opt_numpy_obs_prob))
        self.assert_((Pi==self.hmm._opt_numpy_start_prob).all()                , 
                'Pi%s\n---\nPi!%s'%(Pi,self.hmm._opt_numpy_start_prob))
    
    def testOptimized_NoNumpy(self):
        self.hmm._numpy = False
        self.hmm._optimize()
        self.assertAlmostEqual(self.hmm.forward(self.obs_sequence),0.00004374489737513661384582519531)
        self.assertEquals(self.hmm.viterbi(self.obs_sequence),self.result)
    
    def testOptimized_Numpy(self):
        self.hmm._numpy = True
        self.hmm._optimize()
        self.assertAlmostEqual(self.hmm.forward(self.obs_sequence),0.00004374489737513661384582519531)
        self.assertEquals(self.hmm.viterbi(self.obs_sequence),self.result)
    
    def testUnoptimized(self):
        self.assertAlmostEqual(self.hmm.forward(self.obs_sequence),0.00004374489737513661384582519531)
        self.assertEquals(self.hmm.viterbi(self.obs_sequence),self.result)
    
    def testOptimized_Numpy_fast(self):
        self.hmm._numpy = True
        self.hmm._optimize()
        self.assertAlmostEqual(self.hmm.forward(self.obs_sequence),0.00004374489737513661384582519531)
        self.assertEquals(self.hmm._opt_viterbi_numpy(self.obs_sequence),self.result)
        
class HMMTestCase2(unittest.TestCase):
    def setUp(self):
        state_transitions = {
         'sun':{'sun':0.5, 'cloud':0.25, 'rain':0.25}, 
         'cloud':{'sun':0.275, 'cloud':0.125, 'rain':0.375, 'storm':0.1}, 
         'rain':{'sun':0.125, 'cloud':0.475, 'rain':0.375, 'storm':0.2},
         'storm':{'cloud':0.3, 'rain':0.7}
                         }     
        start_probabilities = {'sun':0.63, 'cloud':0.17, 'rain':0.20,}

        states = ['sun', 'cloud', 'rain', 'storm']
        observations = ['dry', 'dryish', 'damp', 'soggy']

        observations_prob = {
        'sun':{'dry':0.60, 'dryish':0.20, 'damp':0.15, 'soggy':0.05}, 
        'cloud':{'dry':0.25, 'dryish':0.25, 'damp':0.25, 'soggy':0.25}, 
        'rain':{'dry':0.05, 'dryish':0.10, 'damp':0.35, 'soggy':0.50},
        'storm':{'damp':0.1, 'soggy':0.9}
        }

        self.hmm = MyHMM(states, observations, start_probabilities, state_transitions, observations_prob)
        self.hmm._calc_coef()
        self.obs_sequence = ['dry', 'damp', 'soggy', 'soggy','damp','damp','dry','dry']
        self.result = ['sun', 'rain', 'rain', 'storm', 'rain', 'cloud', 'sun', 'sun']
    
    def tearDown(self):
        """docstring for tearDown"""
        pass
        
    def testOptimizeMethod(self):
        import numpy
        A = numpy.array([
        [0.5, 0.25, 0.25, self.hmm.coef_trans],
        [0.275,0.125,0.375,0.1],
        [0.125,0.475,0.375, 0.2],
        [self.hmm.coef_trans, 0.3, 0.7, self.hmm.coef_trans]
        ])
        Pi = numpy.array(
        [0.63,0.17,0.20, self.hmm.coef_trans]
        )
        B = numpy.array([
        [0.60,0.20,0.15,0.05],
        [0.25,0.25,0.25,0.25],
        [0.05,0.10,0.35,0.50],
        [self.hmm.coef_obs, self.hmm.coef_obs, 0.1, 0.9]
        ])
        self.hmm._numpy = True
        self.hmm._optimize()
        # start_probabilities = self._opt_numpy_start_prob
        # observations_prob = self._opt_numpy_obs_prob
        # state_transitions = self._opt_numpy_state_trans
        # print 'A%s\n---\nA!%s'%(A,self.hmm._opt_numpy_state_trans)
        # self.assert_(A.all()==self.hmm._opt_numpy_obs_prob.all())
        self.assert_((A==self.hmm._opt_numpy_state_trans).all(), 
                'A%s\n---\nA!%s'%(A,self.hmm._opt_numpy_state_trans))
        self.assert_((B==self.hmm._opt_numpy_obs_prob).all(), 
                'B%s\n---\nB!%s'%(B,self.hmm._opt_numpy_obs_prob))
        self.assert_((Pi==self.hmm._opt_numpy_start_prob).all() , 
                'Pi%s\n---\nPi!%s'%(Pi,self.hmm._opt_numpy_start_prob))
                
    def testUnoptimized(self):
        self.assertAlmostEqual(self.hmm.forward(self.obs_sequence),7.4654138163903868e-05)
        self.assertEquals(self.hmm.viterbi(self.obs_sequence),self.result)
        
    def testOptimized_NoNumpy(self):
        self.hmm._numpy = False
        self.hmm._optimize()
        self.assertAlmostEqual(self.hmm.forward(self.obs_sequence),7.4654138163903868e-05)
        self.assertEquals(self.hmm._opt_viterbi(self.obs_sequence),self.result)
        
    def testOptimized_Numpy(self):
        self.hmm._numpy = True
        self.hmm._optimize()
        self.assertAlmostEqual(self.hmm.forward(self.obs_sequence),7.4654138163903868e-05)
        self.assertEquals(self.hmm._opt_viterbi_numpy2(self.obs_sequence),self.result)    
        
    def testOptimized_Numpy_fast(self):
        self.hmm._numpy = True
        self.hmm._optimize()
        self.assertAlmostEqual(self.hmm.forward(self.obs_sequence),7.4654138163903868e-05)
        self.assertEquals(self.hmm._opt_viterbi_numpy(self.obs_sequence),self.result)

if __name__ == '__main__':
    unittest.main()

